import random
import time

import gym
import numpy as np
from gym.spaces import Discrete

from q_table.Games.GameMaze import GameMaze
from q_table.Games.MazeDTO import generate_maze


class SnakeEnv(gym.Env):
    SIZE = 100

    def __init__(self, ladder_num, dices):
        self.ladder_num = ladder_num
        self.dices = dices
        self.observation_space = Discrete(self.SIZE + 1)
        self.action_space = Discrete(len(dices))

        if ladder_num == 0:
            self.ladders = {0: 0}
        else:
            # 处理梯子值，让梯子的数值无重复地反向赋值
            ladders = set(np.random.randint(1, self.SIZE, size=self.ladder_num * 2))
            while len(ladders) < self.ladder_num * 2:
                ladders.add(np.random.randint(1, self.SIZE))

            ladders = list(ladders)
            ladders = np.array(ladders)
            np.random.shuffle(ladders)
            ladders = ladders.reshape((self.ladder_num, 2))

            re_ladders = list()
            for i in ladders:
                re_ladders.append([i[1], i[0]])

            re_ladders = np.array(re_ladders)
            # dict()可以把nx2维数组转化为字典形式
            self.ladders = dict(np.append(re_ladders, ladders, axis=0))
        print(f'ladders info:{self.ladders} dice ranges:{self.dices}')
        self.pos = 1

    def reset(self, *args, **kwargs):
        self.pos = 1
        return self.pos

    def step(self, a):
        for i in range(a + 2):
            self.pos += np.random.randint(1, 7)
        if self.pos == 100:
            return 100, 100, 1, {}
        elif self.pos > 100:
            self.pos = 200 - self.pos

        if self.pos in self.ladders:
            self.pos = self.ladders[self.pos]
        return self.pos, -1, 0, {}

    def reward(self, s):
        if s == 100:
            return 100
        else:
            return -1
